use std::path::PathBuf;

use anyhow::{Result, bail};
use chrono::{DateTime, Utc};
use tensorzero_core::inference::types::ContentBlockChatOutput;
use tree_sitter::Tree;
use uuid::{Builder, Uuid};

use crate::{clickhouse::InferenceInfo, ted::TedInfo};

#[derive(Debug)]
pub struct TreeInfo {
    pub path: PathBuf, // VSCode workspace relative path for inferences, git-relative path for diffs
    pub tree: Tree,
    pub src: Vec<u8>,
}

#[derive(Debug)]
pub struct NormalizedInferenceTreeInfo {
    pub paths: Vec<PathBuf>, // git-relative paths that might be the right path for this inference
    pub tree: Tree,
    pub src: Vec<u8>,
}

/// Generates the maximum UUIDv7 for a given timestamp.
pub fn get_max_uuidv7(timestamp: DateTime<Utc>) -> Uuid {
    // Create a byte array of 255s
    let bytes: [u8; 10] = [255; 10];
    // Get the unix timestamp in milliseconds as a u64
    let timestamp_ms = timestamp.timestamp_millis() as u64;
    // Create a builder
    let builder = Builder::from_unix_timestamp_millis(timestamp_ms, &bytes);
    // Build the UUID
    builder.into_uuid()
}

/// Generates the minimum UUIDv7 for a given timestamp.
pub fn get_min_uuidv7(timestamp: DateTime<Utc>) -> Uuid {
    // Create a byte array of 0s
    let bytes: [u8; 10] = [0; 10];
    // Get the unix timestamp in milliseconds as a u64
    let timestamp_ms = timestamp.timestamp_millis() as u64;
    // Create a builder
    let builder = Builder::from_unix_timestamp_millis(timestamp_ms, &bytes);
    // Build the UUID
    builder.into_uuid()
}

/// Generates a demonstration using the code blocks found by the search for minimum TED for each code block
/// generated by the inference.
/// We substitute the original code block with the minimum TED source code block.
pub fn generate_demonstration(
    original_inference: &InferenceInfo,
    min_teds: &[TedInfo],
    inference_trees: &[&TreeInfo],
) -> Result<String> {
    if min_teds.len() != inference_trees.len() {
        bail!(
            "min_teds and inference_trees have different lengths: {} != {}",
            min_teds.len(),
            inference_trees.len()
        );
    }
    let mut output_text = match original_inference.output.as_slice() {
        [ContentBlockChatOutput::Text(t)] => t.text.clone(),
        _ => {
            bail!("Output is not a single text block");
        }
    };

    for (min_ted, tree_info) in min_teds.iter().zip(inference_trees.iter()) {
        let original_text = String::from_utf8(tree_info.src.clone())?;
        let Some(min_ted_source) = &min_ted.min_ted_source else {
            // We still want to send the demonstration since there may be other good code blocks.
            tracing::warn!("No min_ted_source for min_ted: {min_ted:?}");
            continue;
        };
        // TODO: this does a clone every time, which is inefficient.
        output_text = output_text.replace(&original_text, min_ted_source);
    }
    Ok(output_text)
}
